#pragma once

#include <arpa/inet.h>
#include <assert.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>

#include <atomic>
#include <deque>
#include <future>
#include <map>
#include <stdexcept>
#include <boost/lockfree/spsc_queue.hpp>

#include "log.h"
#include "config.h"
#include "protocol.h"
#include "rdma.h"

// RDMA send buffer
// because write_cache will be invoked asynchronously,
// so each request will have a standalone send buffer.
struct SendBuffer {
    SendBuffer(const SendBuffer &) = delete;

    /**
     * @brief 分配指定大小的内存空间并注册为本地写MR
     * 
     * @param pd    MR所属保护域
     * @param size  内存空间大小(未使用)
     */
    SendBuffer(struct ibv_pd *pd, size_t size);
    ~SendBuffer();
    
    void    *buffer_   = nullptr;
    struct ibv_mr *mr_ = nullptr;
};




enum class WrType {
    BASE,
    RDMA_READ_ACK,
    RDMA_WRITE_ACK,
};




struct rdma_info_base {
public:
    rdma_info_base(WrType wr_type) : wr_type(wr_type) {}

    virtual ~rdma_info_base() = default;

    WrType get_wr_type() const { return wr_type; }

protected:
    WrType wr_type;
};




struct rdma_write_info : rdma_info_base {
    using Callback = std::function<void(int)>;

    Callback callback;
    rdma_write_info(Callback callback)
        : rdma_info_base(WrType::RDMA_WRITE_ACK), 
            callback(callback) {}
};




struct rdma_read_info : rdma_info_base {
    using Callback = std::function<void(unsigned int)>;

    Callback callback;
    rdma_read_info(Callback callback)
        : rdma_info_base(WrType::RDMA_READ_ACK), 
        callback(callback) {}
};




class Connection {
public:
    Connection() = default;

    Connection(const Connection &) = delete;   // 禁用拷贝构造函数


    /**
     *! @brief 析构函数, 释放所有内存空间, 取消MR，关闭RDMA设备
     */
    ~Connection();

public:
    // close cq_handler thread
    /**
     *! @brief 通过以太网套接字链接到config指定的服务器
     * 
     * @param config 其中包含对端TCP IP + Port
     * @return int   返回连接状态码
     */
    int init_conn(client_config_t config);



    /**
     * @brief 
     * 
     */
    void close_conn();



    /**
     *! @brief 打开本地设备，通过TCP套接字交换RDMA连接信息并设置QP状态到RTS
     * 
     * @param config 
     * @return int 
     */
    int setup_rdma(client_config_t config);




    /**
     *! @brief 使用套接字形式将本地RDMA连接信息发送
     *        给对端并接受对端RDMA连接信息
     * 
     *        填充local_info_和remote_info_内容
     * 
     * @return int 
     */
    int exchange_conn_info();



    /**
     *! @brief 检验InfiniStore中是否存在key(TCP套接字)
     * 
     * @param key   待检验存在性的key
     * @return int  0:不存在 1:存在
     */
    int check_exist(std::string key);




    /**
     *! @brief  通过RDMA READ将数据从InfiniStore读取到本地
     * 
     * @param keys        key列表
     * @param offsets     每个key数据存放的地址偏移
     * @param block_size  块大小
     * @param base_ptr    数据基础地址指针
     * @param callback    回调函数
     * @return int        写入成功状态
     */
    int r_rdma_async(const std::vector<std::string> &keys, 
                    const std::vector<size_t> offsets,
                    int block_size, void *base_ptr, 
                    std::function<void(unsigned int)> callback);
    
    
    

    /**
     *! @brief  通过RDMA Write将数据写入InfiniStore
     * 
     * @param keys        key列表
     * @param offsets     每个key数据存放的地址偏移
     * @param block_size  块大小
     * @param base_ptr    数据基础地址指针
     * @param callback    回调函数
     * @return int        写入成功状态
     */
    int w_rdma_async(const std::vector<std::string> &keys, 
                    const std::vector<size_t> offsets,
                    int block_size, void *base_ptr, 
                    std::function<void(int)> callback);
    
    
    
    
    /**
     *! @brief 通过TCP套接字将数据写入InfiniStore
     * 
     * @param key   数据对应的key
     * @param ptr   数据存放的地址
     * @param size  数据大小
     * @return int  写入状态
     */
    int w_tcp(const std::string &key, void *ptr, size_t size);
    
    
    

    /**
     *! @brief 通过TCP套接字读取InfiniStore中的数据
     * 
     * @param key 要读取的KV Cache的健
     * @return std::vector<unsigned char>* 返回读取的数据地址
     */
    std::vector<unsigned char> *r_tcp(const std::string &key);

    
    
    
    int get_match_last_index(std::vector<std::string> &keys);
    
    

    
    /**
     *! @brief  删除指定健的KV数据
     * 
     * @param keys  键列表
     * @return int  返回成功删除的key对应的KV数据的数量
     */
    int delete_keys(const std::vector<std::string> &keys);
    
    
    
    
    /**
     *! @brief 普通的MR注册
     * 
     * @param ptr 
     * @param size 
     * @return int 
     */
    int register_mr(void *ptr, size_t size);

 


    /**
     *! @brief 下发一个接收ACK信息的Recv WR
     * 
     * @param info info中包含WR的worker id
     */
    void post_recv_ack(rdma_info_base *info);

    
    
    
    /**
     *! @brief IP完成事件处理回调函数
     * 
     */
    void cq_handler();
    
    
    
    

    /**
     *! @brief 从Sendbuffer队列中获取一个SendBuffer
     * 
     * @return SendBuffer* 返回获取的SendBuffer对象指针
     */
    SendBuffer *get_send_buffer();
    
    

    
    /**
     *! @brief 将SendBuffer对象放回Sendbuffer队列
     * 
     * @param buffer SendBuffer对象指针
     */
    void release_send_buffer(SendBuffer *buffer);

    
    
    
    //! 原文此两方法未实现
    SendBuffer *get_recv_buffer();
    void release_recv_buffer(SendBuffer *buffer);

private:
    // tcp socket
    int sock_ = 0;

    struct rdma_device  rdma_dev_;
    struct rdma_context ctx_;

    rdma_conn_info_t local_info_;
    rdma_conn_info_t remote_info_;

    // uintptr_t = unsigned long int 
    std::unordered_map<uintptr_t, struct ibv_mr *> local_mr_;

    /*
    This is MAX_RECV_WR not MAX_SEND_WR,
    because server also has the same number of buffers
    */

    // struct ibv_comp_channel *comp_channel_ = NULL;
    std::future<void> cq_future_;  // cq thread

    std::atomic<bool> stop_{false};


    // MAX_RECV_WR = 128
    //! boost::lockfree::spsc_queue 是 Boost 库提供的一个高性能单生产者单消费者
    //（Single Producer Single Consumer, SPSC）无锁队列，适用于需要高吞吐、低
    // 延迟的线程间通信场景。以下是其核心特性和使用方法的详细说明：
    //  - 无锁设计	 生产者消费者线程无需互斥锁，通过原子操作实现线程安全
    //  - 固定容量	 必须在构造时指定队列大小，不可动态扩展
    //  - 极低延迟	 适合高频小数据传递（通常比std::queue+互斥锁快10倍以上）
    //  - 内存连续性 数据存储在连续内存中，缓存友好
    //  - 平台支持	 依赖硬件支持的原子操作（x86/ARM等主流架构均支持）
    //! 常用方法； push, pop, empty, size
    boost::lockfree::spsc_queue<SendBuffer *> send_buffers_{MAX_RECV_WR};
};

